# export_to_excel.py
import os
import argparse
from datetime import datetime
import pandas as pd
from db_pool import get_db_connection


def export_data_to_excel(query, output_file=None):
    """
    从数据库查询数据并导出到Excel文件

    Args:
        query (str): SQL查询语句
        output_file (str, optional): 输出Excel文件名. 默认为None，会自动生成

    Returns:
        str: 生成的Excel文件路径
    """
    # 如果没有指定输出文件名，自动生成
    if not output_file:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_file = f"export_{timestamp}.xlsx"

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_file) if os.path.dirname(output_file) else ".", exist_ok=True)

    connection = None
    cursor = None
    try:
        # 从连接池获取连接
        connection = get_db_connection()
        cursor = connection.cursor()

        # 执行查询
        cursor.execute(query)

        # 获取列名
        columns = [desc[0] for desc in cursor.description]

        # 获取所有数据
        data = cursor.fetchall()

        # 创建DataFrame
        df = pd.DataFrame(data, columns=columns)

        # 导出到Excel
        df.to_excel(output_file, index=False, engine='openpyxl')

        print(f"数据已成功导出到: {output_file}")
        print(f"导出记录数: {len(df)}")

        return output_file

    except Exception as e:
        print(f"导出数据时发生错误: {str(e)}")
        raise
    finally:
        # 确保游标和连接被关闭
        if cursor:
            cursor.close()
        if connection:
            connection.close()


def main():
    """主函数，处理命令行参数"""
    query = "select * from youd.z_products_g"
    # 执行导出
    export_data_to_excel(query)


if __name__ == "__main__":
    main()